#!/usr/bin/env python
# -*- encoding: utf-8 -*-
"""
@File    :   args.py
@Contact :   ruaqy@qq.com
@License :   (C)Copyright 2022-~, GPL 2.0

@Modify Time      @Author    @Version    @Description
------------      -------    --------    -----------
2022/5/16 18:44   rqy        1.0         None
"""
import argparse
from model.info import Info


def parser_init():
    parser = argparse.ArgumentParser(description="ResNet Train Toolkits.\nAuthor:RQY, LICENSE:GPL-2.",
                                     formatter_class=argparse.ArgumentDefaultsHelpFormatter)

    # 版本信息
    prog_info = Info()
    parser.add_argument('-v', '--version', action='version',
                        version='Version {}'.format(prog_info.version))

    # 项目设置
    parser.add_argument('--project', type=str, required=True,
                        help='the name of project.')
    parser.add_argument('--target', type=str, required=True,
                        default='mindspore',
                        choices=['mindspore', 'pytorch'],
                        help='the target device to run training.')

    # 路径参数
    parser.add_argument('--data_url', type=str, required=True,
                        help='the url of train data.')
    parser.add_argument('--save_url', type=str, required=True,
                        help='the url to save model.')

    # 网络模型配置
    parser.add_argument('--net', type=str, default='resnet18',
                        choices=['resnet18', 'resnet50'],
                        help='the network to train.')
    parser.add_argument('--class_num', type=int, default=10,
                        help='The number of classification')
    parser.add_argument('--input_size', type=int, default=224,
                        help='The size of input')
    parser.add_argument('--optim', type=str, default='momentum',
                        choices=['adam', 'momentum'],
                        help='the optimizer choice.')
    parser.add_argument('--callback', type=str, action='append',
                        choices=['LossMonitor', 'ModelCheckpoint', 'SummaryCollector', 'AccuracyMonitor'],
                        help='the callback function in train.')
    parser.add_argument('--freeze', type=str,
                        help='freeze the Conv layer.')

    # 训练参数配置
    parser.add_argument('--epochs', type=int, default=1,
                        help='The epoch of train.')
    parser.add_argument('--batch', type=int, default=1,
                        help='The batch size of train.')

    # 超参数配置
    parser.add_argument('--lr', type=float, default=1e-3,
                        help='the learning rate of optimizer.')
    parser.add_argument('--mm', type=float, default=0.99,
                        help='the momentum rate of optimizer.')
    parser.add_argument('--b1', type=float, default=0.9,
                        help='the first  momentum decay rate of optimizer.')
    parser.add_argument('--b2', type=float, default=0.999,
                        help='the second momentum decay rate of optimizer.')
    parser.add_argument('--wd', type=float, default=0.0,
                        help='the weight  decay of optimizer.')

    # 训练数据配置
    parser.add_argument('--unzip', type=str,
                        help='the unzip file relative to data_url.')
    parser.add_argument('--unzip_target', type=str, default='./',
                        help='the unzip file target path.')

    parser.add_argument('--image_dir', type=str, default='images/',
                        help='the image path relative to data_url.')
    parser.add_argument('--label_dir', type=str, default='dataset.txt',
                        help='the label path relative to data_url.')

    parser.add_argument('--pretrain', type=str, default='resnet',
                        help='pretrain model path relative to data_url.')

    parser.add_argument('--export', type=str, default='resnet',
                        help='the file name of export networks.')

    return parser


if __name__ == '__main__':
    parser_init()
